# coding=utf-8
import argparse
import pickle
import pprint
import random
import sys
import string

parser = argparse.ArgumentParser()
parser.add_argument('-m', action='store', dest='method',
                    help='操作方式 w：写密码， r：查找密码， g：生成密码')
parser.add_argument('-c', action='store', dest=' category',
                    help='类别')
parser.add_argument('-i', action='store', dest='title',
                    help='查找的名称')
parser.add_argument('-l', type=int, action='store', dest='length',
                    help='密码长度')
args = parser.parse_args()


'''
密码管理器

* 保存密码
* 获取密码

格式：
 * 密码 title
 * {"title": "passwd"}

分类：
 * web
 * 服务器

格式：
 {
 "web":{"jingfeng":"123456", "zhihu":"123456"}
 }
 title | passwd | url

自动生成安全密码 （用户可以输入长度）
'''


def select_db():
    try:
        with open('zpass.pass', 'rb') as f:
            return pickle.load(f)
    except Exception as e:
        return False



def insert_db(obj):
    with open('zpass.pass', 'wb') as f:
        pickle.dump(obj, f)


class PassWord(dict):
    
    def __init__(self, category, title, user, passwd):
        '''
        {'title':
                    {
                        user:test, 
                        passwd:123456, 
                        url:jingfeng.com
                    }
        }
        '''
        super().__init__()
        self.title = title
        self.category = category
        self[title] = {'user': user, 'passwd': passwd}


    def passwd_data(self):
        '''
        {'title': 'test', 'user': 'test', 'passwd': '123456', 'url': 'url.com'}
        |
        V
        {
            'web':{'title':
                    {
                        user:test, 
                        passwd:123456, 
                        url:jingfeng.com
                    }
                }
        }
        '''
        pwd_data = select_db() if select_db() else {}
        if not pwd_data or self.category not in pwd_data:
            pwd_data[self.category] = {}
            pwd_data[self.category][self.title] = self[self.title]
        pwd_data[self.category][self.title] = self[self.title]
        return pwd_data

    def save_passwd(self):
        passwd_data = self.passwd_data()
        pprint.pprint(passwd_data)
        insert_db(passwd_data)
        return True

    @classmethod
    def get_item(cls):
        cls.item = ['title', 'user', 'passwd']
        return cls.item

    @staticmethod
    def get_passwd():
        passwd =  select_db()
        return passwd[args.category].get(args.title)


class WebPassWord(PassWord):

    def __init__(self, category, title, user, passwd, url):
        super().__init__(category, title, user, passwd)
        self[title]['url'] = url

    @classmethod
    def get_item(cls):
        super().get_item()
        cls.item.append('url')
        return cls.item


class ServerPassWord(PassWord):
    
    def __init__(self, category, title, user, passwd, hostname, ip):
        super().__init__(category, title, user, passwd)
        self[title]['hostname'] = hostname
        self[title]['ip'] = ip

    @classmethod
    def get_item(cls):
        super().get_item()
        cls.item.append('hostname')
        cls.item.append('ip')
        return cls.item


def handle_passwd():
    passwd_args = {}
    user_choice = input("1 web, 2 server, 0 退出")
    pwd_class = {
    '1': ('web', WebPassWord), 
    '2': ('server', ServerPassWord)}
    category = pwd_class[user_choice][0]
    handle_class = pwd_class[user_choice][1]
    for i in handle_class.get_item():
        passwd_args[i] = input(i + ":")
    pwd = handle_class(category=category, **passwd_args)
    stat = pwd.save_passwd()
    if stat is False:
        print('保存失败！')
    print('保存成功！')


def gen_passwd(length):
    chars = string.ascii_letters + string.digits + "!@#$%^&*():;"'`'
    return ''.join([random.choice(chars) for i in range(length)])


def main():
    if args.method == 'r':
        print(PassWord.get_passwd())
    elif args.method == 'w':
        handle_passwd()
    elif args.method == 'g':
        print(gen_passwd(args.length))



if __name__ == '__main__':
    main()
